# -*- coding: utf-8 -*-
"""
    @project: pythonProject
    @Author：HanYonghua
    @file： sweepfreq.py
    @date：2025/4/7 17:21
    @blogs: https://www.ncatest.com.cn
"""
import numpy as np
from scipy.signal import butter, filtfilt
import matplotlib.pyplot as plt

# 生成含噪CW信号示例
fs = 1000  # 采样率 (Hz)
t = np.arange(0, 1, 1 / fs)  # 1秒时间轴
f0_true = 47.3  # 真实信号频率 (Hz)，假设未知
signal = 0.5 * np.sin(2 * np.pi * f0_true * t)  # 原始信号
noisy_signal = signal + 0.8 * np.random.randn(len(t))  # 强噪声淹没


# 扫频法检测频率
def frequency_sweep(signal, fs, f_min, f_max, f_step):
    """扫频法检测信号频率"""
    freqs = np.arange(f_min, f_max, f_step)
    energies = []

    for f in freqs:
        # 生成参考正弦和余弦波
        ref_sin = np.sin(2 * np.pi * f * t)
        ref_cos = np.cos(2 * np.pi * f * t)

        # 计算相干能量（类似FFT的单个频点）
        energy_sin = np.abs(np.sum(signal * ref_sin)) ** 2
        energy_cos = np.abs(np.sum(signal * ref_cos)) ** 2
        total_energy = energy_sin + energy_cos
        energies.append(total_energy)

    # 找到能量最大的频率
    peak_idx = np.argmax(energies)
    f_estimated = freqs[peak_idx]
    return f_estimated, freqs, energies


# 参数设置
f_min = 1  # 最小扫描频率 (Hz)
f_max = fs / 2  # 最大扫描频率 (Nyquist频率)
f_step = 0.1  # 扫描步长 (Hz)

# 执行扫频
f_estimated, freqs, energies = frequency_sweep(noisy_signal, fs, f_min, f_max, f_step)
print(f"真实频率: {f0_true:.2f} Hz | 估计频率: {f_estimated:.2f} Hz")


# 设计带通滤波器恢复信号
def butter_bandpass(lowcut, highcut, fs, order=5):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a


b, a = butter_bandpass(f_estimated - 1, f_estimated + 1, fs)  # 带宽±1Hz
filtered_signal = filtfilt(b, a, noisy_signal)

# 绘制结果
plt.figure(figsize=(14, 10))

# 1. 原始信号与含噪信号
plt.subplot(3, 1, 1)
plt.plot(t, signal, label='Original CW', linewidth=2)
plt.plot(t, noisy_signal, label='Noisy Signal', alpha=0.5)
plt.xlabel('Time (s)')
plt.ylabel('Amplitude')
plt.title('Original vs Noisy Signal')
plt.legend()

# 2. 扫频能量谱
plt.subplot(3, 1, 2)
plt.plot(freqs, energies, label='Sweep Energy')
plt.axvline(f0_true, color='r', linestyle='--', label=f'True Freq: {f0_true} Hz')
plt.axvline(f_estimated, color='g', linestyle=':', label=f'Estimated Freq: {f_estimated:.2f} Hz')
plt.xlabel('Frequency (Hz)')
plt.ylabel('Coherent Energy')
plt.title('Frequency Sweep Detection')
plt.legend()

# 3. 恢复信号对比
plt.subplot(3, 1, 3)
plt.plot(t, signal, label='Original CW', linewidth=2)
plt.plot(t, filtered_signal, label='Recovered CW', alpha=0.8)
plt.xlabel('Time (s)')
plt.ylabel('Amplitude')
plt.title('Recovered Signal vs Original')
plt.legend()

plt.tight_layout()
plt.show()